from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical

from tensorflow.keras import layers
from tensorflow.keras.models import Model,load_model
from tensorflow.keras import optimizers
from tensorflow.keras.losses import categorical_crossentropy






(x_train,y_train),(x_test,y_test) = mnist.load_data()
X_train = x_train.reshape((60000,28,28,1))/255.
X_test = x_test.reshape((10000,28,28,1))/255.

y_train = to_categorical(y_train,num_classes=10)
y_test = to_categorical(y_test,num_classes=10)

tempetature = 3

inputs = layers.Input((28,28,1))
x = layers.Conv2D(16,3)(inputs)
x = layers.BatchNormalization(center=True,scale=False)(x)
x = layers.Activation("relu")(x)
x = layers.Conv2D(16,3,strides=2)(x)
x = layers.BatchNormalization(center=True,scale=False)(x)
x = layers.Activation('relu')(x)
x = layers.Conv2D(32,5)(x)
x = layers.BatchNormalization(center=True,scale=False)(x)
x = layers.Activation('relu')(x)
x = layers.Conv2D(32,5,strides=2)(x)
x = layers.BatchNormalization(center=True,scale=False)(x)
x = layers.Activation('relu')(x)
x = layers.Flatten()(x)
x = layers.Dense(60)(x)
x = layers.BatchNormalization(center=True,scale=False)(x)
x = layers.Activation("relu")(x)
x = layers.Dropout(0.3)(x)
x = layers.Dense(10,activation='softmax')(x)
x = layers.Lambda(lambda t:t/tempetature)(x)
student = Model(inputs,x)
student.compile(optimizer=optimizers.SGD(momentum=0.9,nesterov=True),loss=categorical_crossentropy,metrics=['accuracy'])

teacher = load_model("teacher_model")

student.fit(X_train,teacher.predict(X_train)/tempetature,batch_size=128,epochs=5)
student.evaluate(X_test,y_test/tempetature)
student.save("student_model")

